package com.stress.testing.instrumentation;

import com.stress.testing.common.interceptor.Interceptor;
import com.stress.testing.common.interceptor.InterceptorClassLoader;
import com.stress.testing.common.interceptor.InterceptorClassLoaderHolder;
import com.stress.testing.common.matcher.MethodMatcher;
import com.stress.testing.common.plugin.PluginContext;
import net.bytebuddy.agent.builder.AgentBuilder;
import net.bytebuddy.description.method.MethodDescription;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.dynamic.DynamicType;
import net.bytebuddy.implementation.MethodDelegation;
import net.bytebuddy.matcher.ElementMatcher;
import net.bytebuddy.utility.JavaModule;

import java.net.URL;

public class SimpleTransformer implements AgentBuilder.Transformer {
    private PluginContext pluginContext;

    public SimpleTransformer(PluginContext pluginContext) {
        this.pluginContext = pluginContext;
    }

    @Override
    public DynamicType.Builder<?> transform(DynamicType.Builder<?> builder, TypeDescription typeDescription,
                                            ClassLoader classLoader, JavaModule module) {

        for (MethodMatcher methodMatcher : pluginContext.getPlugin().getClassMatcher().getMethodMatchers()) {
            String interceptor = methodMatcher.getInterceptor();
            // 当一个jar中有多个plugin时，不用重复创建ClassLoader
            InterceptorClassLoader interceptorClassLoader =
                    InterceptorClassLoaderHolder.getInterceptorClassLoader(pluginContext.getPluginURL());
            if (interceptorClassLoader == null) {
                interceptorClassLoader =
                        new InterceptorClassLoader(new URL[]{pluginContext.getPluginURL()}, classLoader);
                InterceptorClassLoaderHolder.setInterceptorClassLoader(pluginContext.getPluginURL(), interceptorClassLoader);
            }

            Interceptor intp = null;
            try {
                Class<?> interceptorClazz = interceptorClassLoader.loadClass(interceptor);
                intp = Interceptor.class.cast(interceptorClazz.newInstance());
                intp.init(pluginContext.getPlugin().getAgentConfig(), pluginContext.getPlugin().getPluginConfig());
            } catch (Exception e) {
                throw new RuntimeException("插件类创建失败：" + interceptor, e);
            }

            ElementMatcher<MethodDescription> simpleMethodElementMatcher = new SimpleMethodElementMatcher(methodMatcher);
            SimpleMethodDelegation delegation = new SimpleMethodDelegation(intp);
            MethodDelegation methodDelegation = MethodDelegation.to(delegation);
            builder = builder.method(simpleMethodElementMatcher).intercept(methodDelegation);
        }

        return builder;
    }
}
